import numpy as np
import torch
from config import *

def build_grid_graph(grid_size=ENV_SIZE):
    """构建网格图的节点和边索引"""
    # 为每个单元格分配一个唯一的节点索引
    node_indices = np.arange(grid_size * grid_size).reshape(grid_size, grid_size)
    
    # 创建边列表
    edge_index = []
    edge_attr = []
    
    # 遍历所有单元格
    for i in range(grid_size):
        for j in range(grid_size):
            node_idx = node_indices[i, j]
            
            # 添加到相邻单元格的边
            for d_idx, (di, dj) in enumerate(DIRECTIONS):
                ni, nj = i + di, j + dj
                
                # 确保相邻单元格在网格内
                if 0 <= ni < grid_size and 0 <= nj < grid_size:
                    neighbor_idx = node_indices[ni, nj]
                    
                    # 添加有向边
                    edge_index.append([node_idx, neighbor_idx])
                    
                    # 边属性：方向的one-hot编码
                    edge_attr_vec = [0] * len(DIRECTIONS)
                    edge_attr_vec[d_idx] = 1
                    edge_attr.append(edge_attr_vec)
    
    # 将列表转换为PyTorch张量
    edge_index = torch.tensor(edge_index, dtype=torch.long).t().contiguous()
    edge_attr = torch.tensor(edge_attr, dtype=torch.float)
    
    return edge_index, edge_attr

def grid_to_graph_features(grid_state, distance_maps=None, direction_maps=None):
    """将网格状态转换为图节点特征"""
    # 网格大小
    grid_size = grid_state.shape[0]
    
    # 初始化节点特征
    node_features = []
    
    # 遍历网格
    for i in range(grid_size):
        for j in range(grid_size):
            # 获取单元格状态
            cell_state = grid_state[i, j]
            
            # 创建基本特征（one-hot编码）
            cell_features = [0] * CELL_FEATURES
            cell_features[cell_state] = 1
            
            # 添加位置编码
            pos_i = 2 * (i / (grid_size - 1)) - 1  # 范围从-1到1
            pos_j = 2 * (j / (grid_size - 1)) - 1  # 范围从-1到1
            
            # 如果有距离图，添加距离特征
            dist_features = []
            if distance_maps is not None:
                for dist_map in distance_maps:
                    dist_features.append(dist_map[i, j])
            
            # 如果有方向图，添加方向特征
            dir_features = []
            if direction_maps is not None:
                for dir_map in direction_maps:
                    dir_features.append(dir_map[i, j])
            
            # 组合所有特征
            node_feature = cell_features + [pos_i, pos_j] + dist_features + dir_features
            node_features.append(node_feature)
    
    # 转换为PyTorch张量
    node_features = torch.tensor(node_features, dtype=torch.float)
    
    return node_features

def create_distance_maps(grid_state):
    """创建到关键元素的距离图"""
    grid_size = grid_state.shape[0]
    
    # 初始化距离图
    obstacle_dist = np.ones((grid_size, grid_size)) * grid_size * 2
    goal_dist = np.ones((grid_size, grid_size)) * grid_size * 2
    agent_dist = np.ones((grid_size, grid_size)) * grid_size * 2
    
    # 找到关键元素位置
    obstacle_positions = np.where(grid_state == OBSTACLE)
    goal_position = np.where(grid_state == END)
    agent_position = np.where(grid_state == AGENT)
    
    # 计算到障碍物的距离
    for i in range(len(obstacle_positions[0])):
        oi, oj = obstacle_positions[0][i], obstacle_positions[1][i]
        for i in range(grid_size):
            for j in range(grid_size):
                # 曼哈顿距离
                dist = abs(i - oi) + abs(j - oj)
                obstacle_dist[i, j] = min(obstacle_dist[i, j], dist)
    
    # 计算到目标的距离
    if len(goal_position[0]) > 0:
        gi, gj = goal_position[0][0], goal_position[1][0]
        for i in range(grid_size):
            for j in range(grid_size):
                # 曼哈顿距离
                goal_dist[i, j] = abs(i - gi) + abs(j - gj)
    
    # 计算到智能体的距离
    if len(agent_position[0]) > 0:
        ai, aj = agent_position[0][0], agent_position[1][0]
        for i in range(grid_size):
            for j in range(grid_size):
                # 曼哈顿距离
                agent_dist[i, j] = abs(i - ai) + abs(j - aj)
    
    # 归一化距离图
    max_possible_dist = grid_size * 2
    obstacle_dist = obstacle_dist / max_possible_dist
    goal_dist = goal_dist / max_possible_dist
    agent_dist = agent_dist / max_possible_dist
    
    return [obstacle_dist, goal_dist, agent_dist]

def create_direction_maps(grid_state):
    """创建指向关键元素的方向图"""
    grid_size = grid_state.shape[0]
    
    # 初始化方向图 (sin角度, cos角度)
    goal_dir = np.zeros((grid_size, grid_size, 2))
    agent_dir = np.zeros((grid_size, grid_size, 2))
    
    # 找到关键元素位置
    goal_position = np.where(grid_state == END)
    agent_position = np.where(grid_state == AGENT)
    
    # 计算指向目标的方向
    if len(goal_position[0]) > 0:
        gi, gj = goal_position[0][0], goal_position[1][0]
        for i in range(grid_size):
            for j in range(grid_size):
                if i == gi and j == gj:
                    continue
                    
                # 计算方向向量
                di = gi - i
                dj = gj - j
                # 计算距离
                dist = np.sqrt(di**2 + dj**2)
                if dist > 0:
                    # 归一化并计算sin和cos
                    di /= dist
                    dj /= dist
                    goal_dir[i, j, 0] = di  # sin角度
                    goal_dir[i, j, 1] = dj  # cos角度
    
    # 计算指向智能体的方向
    if len(agent_position[0]) > 0:
        ai, aj = agent_position[0][0], agent_position[1][0]
        for i in range(grid_size):
            for j in range(grid_size):
                if i == ai and j == aj:
                    continue
                    
                # 计算方向向量
                di = ai - i
                dj = aj - j
                # 计算距离
                dist = np.sqrt(di**2 + dj**2)
                if dist > 0:
                    # 归一化并计算sin和cos
                    di /= dist
                    dj /= dist
                    agent_dir[i, j, 0] = di  # sin角度
                    agent_dir[i, j, 1] = dj  # cos角度
    
    # 重塑方向图以便于使用
    goal_dir_flat = [goal_dir[:, :, 0], goal_dir[:, :, 1]]
    agent_dir_flat = [agent_dir[:, :, 0], agent_dir[:, :, 1]]
    
    return goal_dir_flat + agent_dir_flat